/**
 * @file     ring_queue.hpp
 * @author   WMQ
 * @date     2024/11/17
 * @brief    无锁环形队列
 */

#ifndef RING_QUEUE_HPP
#define RING_QUEUE_HPP

#include <atomic>
#include <cstdint>
#include <thread>

/**
 * @brief 无锁环形队列
 */
template <typename T>
class RingQueue
{
public:
    /**
     * @brief 构造函数
     * @param[in] size 队列长度，大于或等于size的最小2的幂指数
     * @param[in] args 队列元素对象的构造参数
     */
    template <typename... Args>
    explicit RingQueue(uint32_t size, Args &&... args)
        : producer_head_(0), producer_tail_(0), consumer_head_(0), consumer_tail_(0), size_(size)
    {
        // 求大于或等于size_的最小2的幂指数
        size_ = PowerOfTwo(size_);
        if (size_ <= 1)
        {
            size_ = RING_QUEUE_DEFAULT_LENGTH;
        }
        mask_ = size_ - 1;
        ring_buff_ = static_cast<T*>(::operator new(sizeof(T) * size_));
        for (uint32_t i = 0; i < size_; ++i)
        {
            new (ring_buff_ + i) T(std::forward<Args>(args)...);
        }
    }

    /**
     * @brief 析构函数
     */
    virtual ~RingQueue()
    {
        for (uint32_t i = 0; i < size_; ++i)
        {
            (ring_buff_ + i)->~T();
        }
        ::operator delete(ring_buff_);
    }

    /**
     * @brief 单生产者，一次可以写入多个，单线程调用
     * @param[in] value_table 入队列元素数组的首地址
     * @param[in] n 入队列元素的个数，不能大于数组的长度
     * @return
     *  false   操作失败 \n
     *  true    操作成功 \n
     */
    bool SingleProducerEnQueue(const T* value_table, uint32_t n)
    {
        if (n < 1)
        {
            return false;
        }

        const uint32_t prod_head = this->producer_head_.load(std::memory_order_acquire);
        const uint32_t cons_tail = this->consumer_tail_.load(std::memory_order_acquire);
        // free_entries 结果总是在0 ~ size(ring)-1 = this->mask_
        const uint32_t free_entries = this->mask_ + cons_tail - prod_head;

        // 判断是否有足够的容量
        if (n > free_entries)
        {
            return false;
        }

        const uint32_t prod_next = prod_head + n;
        this->producer_head_.store(prod_next, std::memory_order_release);

        // 将元素写入队列
        EnQueue(prod_head, value_table, n);

        this->producer_tail_.store(prod_next, std::memory_order_release);
        return true;
    }

    /**
     * @brief 单消费者，一次可以取出多个，单线程调用
     * @param[in] value_table 存储队列元素数组的首地址
     * @param[in] n 取出元素的个数，不能大于数组的长度
     * @return
     *  false   操作失败 \n
     *  true    操作成功 \n
     */
    bool SingleConsumerDeQueue(T* value_table, uint32_t n)
    {
        if (n < 1)
        {
            return false;
        }

        const uint32_t cons_head = this->consumer_head_.load(std::memory_order_acquire);
        const uint32_t prod_tail = this->producer_tail_.load(std::memory_order_acquire);
        // entries 结果总是在0 ~ size(ring)-1 = this->mask_
        const uint32_t entries = prod_tail - cons_head;

        if (n > entries)
        {
            return false;
        }

        const uint32_t cons_next = cons_head + n;
        this->consumer_head_.store(cons_next, std::memory_order_release);

        // 从队列中取出元素
        DeQueue(cons_head, value_table, n);

        this->consumer_tail_.store(cons_next, std::memory_order_release);
        return true;
    }

    /**
     * @brief 多生产者，一次可以写入多个，多线程调用
     * @param[in] value_table 入队列元素数组的首地址
     * @param[in] n 入队列元素的个数，不能大于数组的长度
     * @return
     *  false   操作失败 \n
     *  true    操作成功 \n
     */
    bool MultiProducerEnQueue(const T* value_table, uint32_t n)
    {
        if (n < 1)
        {
            return false;
        }

        uint32_t prod_head = 0;
        uint32_t prod_next = 0;
        unsigned rep = 0;

        do
        {
            prod_head = this->producer_head_.load(std::memory_order_acquire);
            const uint32_t cons_tail = this->consumer_tail_.load(std::memory_order_acquire);
            // free_entries 结果总是在0 ~ size(ring)-1 = this->mask_
            const uint32_t free_entries = (this->mask_ + cons_tail - prod_head);

            // 判断是否有足够的容量
            if (n > free_entries)
            {
                return false;
            }

            prod_next = prod_head + n;
        } while (!this->producer_head_.compare_exchange_weak(
            prod_head, prod_next, std::memory_order_release, std::memory_order_relaxed));

        // 将元素写入队列
        EnQueue(prod_head, value_table, n);

        // 如果在这之前还有其他线程正在写入，需要等待它们完成
        while (this->producer_tail_.load(std::memory_order_acquire) != prod_head)
        {
            if (++rep == RING_QUEUE_WAIT_COUNT)
            {
                rep = 0;
                // 重新调度线程的执行以允许其他线程运行
                std::this_thread::yield();
            }
        }
        this->producer_tail_.store(prod_next, std::memory_order_release);
        return true;
    }

    /**
     * @brief 多消费者，一次可以取出多个，多线程调用
     * @param[in] value_table 存储队列元素数组的首地址
     * @param[in] n 取出元素的个数，不能大于数组的长度
     * @return
     *  false   操作失败 \n
     *  true    操作成功 \n
     */
    bool MultiConsumerDeQueue(T* value_table, uint32_t n)
    {
        if (n < 1)
        {
            return false;
        }

        uint32_t cons_head = 0;
        uint32_t cons_next = 0;
        unsigned rep = 0;

        do
        {
            cons_head = this->consumer_head_.load(std::memory_order_acquire);
            const uint32_t prod_tail = this->producer_tail_.load(std::memory_order_acquire);
            // entries 结果总是在0 ~ size(ring)-1 = this->mask_
            const uint32_t entries = (prod_tail - cons_head);

            // 判断是否可以取出n个元素
            if (n > entries)
            {
                return false;
            }

            cons_next = cons_head + n;
        } while (!this->consumer_head_.compare_exchange_weak(
            cons_head, cons_next, std::memory_order_release, std::memory_order_relaxed));

        // 从队列中取出元素
        DeQueue(cons_head, value_table, n);

        // 如果在这之前还有其他线程正在取元素，需要等待它们完成
        while (this->consumer_tail_.load(std::memory_order_acquire) != cons_head)
        {
            if (++rep == RING_QUEUE_WAIT_COUNT)
            {
                rep = 0;
                // 重新调度线程的执行以允许其他线程运行
                std::this_thread::yield();
            }
        }
        this->consumer_tail_.store(cons_next, std::memory_order_release);
        return true;
    }

    /**
     * @brief 返回队列最大容量
     */
    uint32_t GetSize()
    {
        return size_ - 1;
    }

private:
    uint32_t size_;                                   /**< 队列长度 */
    uint32_t mask_;                                   /**< 掩码，a % b(2^n) = a & (b - 1) */
    std::atomic<uint32_t> producer_head_ = 0;         /**< 生产者头指针 */
    std::atomic<uint32_t> producer_tail_ = 0;         /**< 生产者尾指针 */
    std::atomic<uint32_t> consumer_head_ = 0;         /**< 消费者头指针 */
    std::atomic<uint32_t> consumer_tail_ = 0;         /**< 消费者尾指针 */
    T* ring_buff_;                                    /**< 队列元素首地址 */
    static const int RING_QUEUE_WAIT_COUNT = 3;       /**< 自旋多少次后执行一次yield() */
    static const int RING_QUEUE_DEFAULT_LENGTH = 128; /**< 默认分配的队列元素个数2^7 */

    /**
     * @brief 求大于或等于一个整数的最小2的幂指数
     * @param[in] num 一个整数
     * @return
     *  0   未找到   \n
     *  >0  所求结果 \n
     */
    uint32_t PowerOfTwo(uint32_t num)
    {
        --num;
        num |= num >> 1;
        num |= num >> 2;
        num |= num >> 4;
        num |= num >> 8;
        num |= num >> 16;
        return ++num;
    }

    /**
     * @brief 从队列取出元素，存入数组
     * @param[in] cons_head 取出第一个元素的位置
     * @param[in] value_table 存储队列元素数组的首地址
     * @param[in] n 取出元素的个数，不能大于数组的长度
     */
    void DeQueue(uint32_t cons_head, T* value_table, uint32_t n)
    {
        uint32_t i = 0;
        uint32_t idx = cons_head & this->mask_;
        const uint32_t size = this->size_;
        if (idx + n < size)
        {
            for (i = 0; i < (n & (~static_cast<uint32_t>(3))); i += 4, idx += 4)
            {
                value_table[i] = std::move(this->ring_buff_[idx]);
                value_table[i + 1] = std::move(this->ring_buff_[idx + 1]);
                value_table[i + 2] = std::move(this->ring_buff_[idx + 2]);
                value_table[i + 3] = std::move(this->ring_buff_[idx + 3]);
            }
            switch (n & static_cast<uint32_t>(3))
            {
            case 3:
                value_table[i++] = std::move(this->ring_buff_[idx++]);
            case 2:
                value_table[i++] = std::move(this->ring_buff_[idx++]);
            case 1:
                value_table[i++] = std::move(this->ring_buff_[idx++]);
            }
        }
        else
        {
            for (i = 0; idx < size; i++, idx++)
            {
                value_table[i] = std::move(this->ring_buff_[idx]);
            }

            for (idx = 0; i < n; i++, idx++)
            {
                value_table[i] = std::move(this->ring_buff_[idx]);
            }
        }
    }

    /**
     * @brief 将数组元素插入队列
     * @param[in] prod_head 插入队列的起始位置
     * @param[in] value_table 入队列元素数组的首地址，第一个入队列元素的地址
     * @param[in] n 入队列元素的个数，不能大于数组的长度
     */
    void EnQueue(uint32_t prod_head, const T* value_table, uint32_t n)
    {
        uint32_t i = 0;
        const uint32_t size = this->size_;
        uint32_t idx = prod_head & this->mask_;
        if (idx + n < size)
        {
            for (i = 0; i < (n & (~static_cast<uint32_t>(3))); i += 4, idx += 4)
            {
                this->ring_buff_[idx] = value_table[i];
                this->ring_buff_[idx + 1] = value_table[i + 1];
                this->ring_buff_[idx + 2] = value_table[i + 2];
                this->ring_buff_[idx + 3] = value_table[i + 3];
            }
            switch (n & static_cast<uint32_t>(3))
            {
            case 3:
                this->ring_buff_[idx++] = value_table[i++];
            case 2:
                this->ring_buff_[idx++] = value_table[i++];
            case 1:
                this->ring_buff_[idx++] = value_table[i++];
            }
        }
        else
        {
            for (i = 0; idx < size; i++, idx++)
            {
                this->ring_buff_[idx] = value_table[i];
            }

            for (idx = 0; i < n; i++, idx++)
            {
                this->ring_buff_[idx] = value_table[i];
            }
        }
    }
};

#endif